from BaseModel import *
import model as model
from collections import OrderedDict
import numpy as np
import soundfile as sf

device = torch.device("cpu")

ckpt = 'aa.ckpt'
model_ = model.DCCRN()
#model_.cuda()

model_info = torch.load(ckpt,map_location=device)
model_info = torch.load(ckpt)
state_dict = OrderedDict()

for k,v in model_info['state_dict'].items():
    name = k.replace("model.", "")
    state_dict[name] = v
model_.load_state_dict(state_dict)
test_data,fs = sf.read('input/item1_source1.wav',dtype=np.float32)
torch_test_data = torch.tensor(test_data,device=device).reshape(1,-1)
test_out = model_(torch_test_data).cpu().detach().reshape(-1, 1).numpy() 
#test_out = model_(torch_test_data).cpu().detach().squeeze().numpy()

test_out = sf.write('out.wav',test_out,fs,'PCM_16')


